package com.archgeek.calcite.test;

import org.apache.calcite.avatica.util.Casing;
import org.apache.calcite.config.Lex;
import org.apache.calcite.config.NullCollation;
import org.apache.calcite.sql.SqlCall;
import org.apache.calcite.sql.SqlDialect;
import org.apache.calcite.sql.SqlIdentifier;
import org.apache.calcite.sql.SqlInsert;
import org.apache.calcite.sql.SqlJoin;
import org.apache.calcite.sql.SqlKind;
import org.apache.calcite.sql.SqlNode;
import org.apache.calcite.sql.SqlSelect;
import org.apache.calcite.sql.dialect.SparkSqlDialect;
import org.apache.calcite.sql.parser.SqlParser;
import org.apache.calcite.sql.validate.SqlConformanceEnum;
import org.apache.calcite.tools.FrameworkConfig;
import org.apache.calcite.tools.Frameworks;
import org.checkerframework.checker.nullness.qual.Nullable;

import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;

/**
 * @author pizhihui
 * @date 2022-04-13 10:29
 */
public class SqlParseDemo {


    public static void main(String[] args) {


        String sql0 = "SELECT MIN(relation_id) FROM tableA JOIN TableB  GROUP BY account_instance_id, follow_account_instance_id HAVING COUNT(*)>1";

        String sql1 = "SELECT * FROM blog_user_relation a WHERE (a.account_instance_id,a.follow_account_instance_id) IN (SELECT account_instance_id,follow_account_instance_id FROM Blogs_info GROUP BY account_instance_id, follow_account_instance_id HAVING COUNT(*) > 1)";
        String sql2 = "select name from (select * from student)";
        String sql3 = "SELECT * FROM Student LEFT JOIN Grade ON Student.sID = Grade.gID\n" +
                "UNION\n" +
                "SELECT * FROM Student RIGHT JOIN Grade ON Student.sID = Grade.gID";
        String sql4 = "SELECT *\n" +
                "FROM teacher\n" +
                "WHERE birth = (SELECT MIN(birth)\n" +
                "               FROM employee)";
        final String sql5 = "SELECT sName\n" +
                "FROM Student\n" +
                "WHERE '450' NOT IN (SELECT courseID\n" +
                "                    FROM Course\n" +
                "                    WHERE sID = Student.sID)";


        String sql6 = "select\n" +
                "    case when ta.`1st_180_month` is not null then 'wo' else 'm0_6' end ftype,\n" +
                "    sum(fpaying_capital)/100 balance\n" +
                "from dp_jckx_mart.mis_union ta \n" +
                "inner join dp_jckx_mart.newml_m3_m6_eco_2022q1_0401 p on ta.forder_id = p.forder_id\n" +
                "left join lx_dws.dws_fin_ord_order_detail_df tb  on ta.forder_id = tb.forder_id\n" +
                "where \n" +
                "    ta.f_p_mon = '2022-03' \n" +
                "    and tb.ffin_one_level_name <> '导流'\n" +
                "group by \n" +
                "    case when ta.`1st_180_month` is not null then 'wo' else 'm0_6' end";

        String sql7 = "drop table dp_ai_mart.ailab_graph_model_call_node_fea_backtrack---------arnoldwang -----";
        String sql8 = "alter table dp_fk_tmp.jy_order_scenes_new_driver_bak1 drop columns subProductCode---------julianzhu -----";

        String sql9 = "create table dp_ai_tmp.second_account_dxzp_Xray_features1_0413 stored as orc as select * from dp_ai_tmp.second_account_dxzp_Xray_features1_0411 union select * from dp_ai_tmp.second_account_dxzp_Xray_features1_0412_1 union select * from dp_ai_tmp.second_account_dxzp_Xray_features1_0412_2 union select * from dp_ai_tmp.second_account_dxzp_Xray_features1_0412_3---------gracetan -----";


        final SqlNode sqlNode0 = parseStatement(sql9);
        System.out.println("sqlNode0: " + extractSourceTableInSelectSql(sqlNode0, false));



            SqlDialect.Context sparkContext = SqlDialect.EMPTY_CONTEXT
                    .withDatabaseProduct(SqlDialect.DatabaseProduct.SPARK)
                    .withNullCollation(NullCollation.LOW)
                    .withCaseSensitive(false)
                    .withConformance(SqlConformanceEnum.BABEL)
                    .withIdentifierQuoteString("`")
                    .withQuotedCasing(Casing.UNCHANGED)
                    .withUnquotedCasing(Casing.UNCHANGED);
        SqlDialect sparkSqlDialect = new SparkSqlDialect(sparkContext);


//        val sqlParser: SqlParser = SqlParser.create(sqlQuery, createSqlParserConfig())
//
//        private def createSqlParserConfig() = { Spark3SqlDialect.DEFAULT
//                .configureParser(SqlParser.config()
//                .withParserFactory(SqlDdlParserImpl.FACTORY))
//                .withConformance(SqlConformanceEnum.BABEL) }

    }

    public static SqlNode parseStatement(String sql) {
        FrameworkConfig config = Frameworks.newConfigBuilder()
                .parserConfig(SqlParser.configBuilder().setLex(Lex.MYSQL).build())
//                .costFactory(costFactory)
//                .typeSystem(typeSystem)
//                .operatorTable(getSqlOperatorTable(calciteConfig(tableConfig), functionCatalog))
//                .sqlToRelConverterConfig(
//                        getSqlToRelConverterConfig(calciteConfig(tableConfig), expressionBridge))
//                // set the executor to evaluate constant expressions
//                .executor(new ExpressionReducer(tableConfig))
                .build();
        SqlParser parser = SqlParser.create(sql, config.getParserConfig());
//        SqlParser parser = SqlParser.create(sql);
        try {
            return parser.parseStmt();
        } catch (Exception e) {
            e.printStackTrace();
            throw new UnsupportedOperationException("operation not allowed");
        }
    }

    private static Set<String> extractSourceTableInSelectSql(SqlNode sqlNode, boolean fromOrJoin) {
        if (sqlNode == null) {
            return new HashSet<>();
        }
        final SqlKind sqlKind = sqlNode.getKind();
        if (SqlKind.SELECT.equals(sqlKind)) {
            SqlSelect selectNode = (SqlSelect) sqlNode;
            Set<String> selectList = new HashSet<>(extractSourceTableInSelectSql(selectNode.getFrom(), true));
            selectNode.getSelectList().getList().stream().filter(node -> node instanceof SqlCall)
                    .forEach(node -> selectList.addAll(extractSourceTableInSelectSql(node, false)));
            selectList.addAll(extractSourceTableInSelectSql(selectNode.getWhere(), false));
            selectList.addAll(extractSourceTableInSelectSql(selectNode.getHaving(), false));
            return selectList;
        }

        if (SqlKind.JOIN.equals(sqlKind)) {
            SqlJoin sqlJoin = (SqlJoin) sqlNode;
            Set<String> joinList = new HashSet<>();
            joinList.addAll(extractSourceTableInSelectSql(sqlJoin.getLeft(), true));
            joinList.addAll(extractSourceTableInSelectSql(sqlJoin.getRight(), true));
            return joinList;
        }
        if (SqlKind.AS.equals(sqlKind)) {
            SqlCall sqlCall = (SqlCall) sqlNode;
            return extractSourceTableInSelectSql(sqlCall.getOperandList().get(0), fromOrJoin);

        }
        if (SqlKind.IDENTIFIER.equals(sqlKind)) {
            Set<String> identifierList = new HashSet<>();
            if (fromOrJoin) {
                SqlIdentifier sqlIdentifier = (SqlIdentifier) sqlNode;
                identifierList.add(sqlIdentifier.toString());
            }
            return identifierList;
        }
        Set<String> defaultList = new HashSet<>();
        if (sqlNode instanceof SqlCall) {
            SqlCall call = (SqlCall) sqlNode;
            call.getOperandList()
                    .forEach(node -> defaultList.addAll(extractSourceTableInSelectSql(node, false)));
            return defaultList;
        }
        return new HashSet<>();
    }

    private static Set<String> extractSourceTableInInsertSql(SqlNode sqlNode, boolean fromOrJoin) {
        SqlInsert sqlInsert = (SqlInsert) sqlNode;
        Set<String> insertList = new HashSet<>(extractSourceTableInSelectSql(sqlInsert.getSource(), false));
        final SqlNode targetTable = sqlInsert.getTargetTable();
        if (targetTable instanceof SqlIdentifier) {
            insertList.add(((SqlIdentifier) targetTable).toString());
        }
        return insertList;
    }
}
